package com.zzl.mq.network.gossip;

import com.alibaba.fastjson.JSON;
import com.zzl.mq.storage.Storage;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.scheduling.annotation.Scheduled;
import org.springframework.stereotype.Component;
import java.util.*;
import java.util.concurrent.ConcurrentHashMap;
import java.util.stream.Collectors;
import javax.annotation.PostConstruct;
import org.springframework.context.annotation.Lazy;

@Component
public class ClusterNodeManager implements GossipMessageHandler {
    private static final Logger log = LoggerFactory.getLogger(ClusterNodeManager.class);
    
    private final Map<String, ClusterNode> nodes = new ConcurrentHashMap<>();
    private final GossipManager gossipManager;
    private final Storage storage;
    
    @Value("${mq.network.heartbeat-timeout}")
    private long heartbeatTimeout;
    
    public ClusterNodeManager(@Lazy GossipManager gossipManager, Storage storage) {
        this.gossipManager = gossipManager;
        this.storage = storage;
    }
    
    @PostConstruct
    public void init() {
        gossipManager.addMessageHandler(this);
    }
    
    @Override
    public void handleMessage(GossipMessage message) {
        switch (message.getType()) {
            case HEARTBEAT:
                updateNodeHeartbeat(message.getNodeId());
                break;
            case NODE_JOIN:
                handleNodeJoin(message);
                break;
            case NODE_LEAVE:
                handleNodeLeave(message);
                break;
            case STATE_UPDATE:
                handleStatusUpdate(message);
                break;
            case DATA_MESSAGE:
                handleDataMessage(message);
                break;
        }
    }
    
    @Scheduled(fixedDelayString = "${mq.network.gossip.cleanup-interval}")
    public void checkNodeHealth() {
        long now = System.currentTimeMillis();
        
        nodes.values().forEach(node -> {
            if (node.getStatus() == NodeStatus.ALIVE 
                && now - node.getLastHeartbeat() > heartbeatTimeout) {
                
                // 标记节点为死亡状态
                node.setStatus(NodeStatus.DEAD);
                
                // 创建节点离开消息
                GossipMessage leaveMessage = new GossipMessage();
                leaveMessage.setMessageId(UUID.randomUUID().toString());
                leaveMessage.setNodeId(node.getNodeId());
                leaveMessage.setType(GossipMessage.MessageType.NODE_LEAVE);
                leaveMessage.setPayload(node);
                leaveMessage.setTimestamp(now);
                
                // 传播节点离开消息
                gossipManager.spreadMessage(leaveMessage);
                
                log.warn("Node {} is marked as DEAD due to heartbeat timeout", node.getNodeId());
            }
        });
    }
    
    // 处理心跳更新
    public void updateNodeHeartbeat(String nodeId) {
        ClusterNode node = nodes.get(nodeId);
        if (node != null) {
            node.updateHeartbeat();
        }
    }
    
    public void addNode(ClusterNode node) {
        nodes.put(node.getNodeId(), node);
    }
    
    public void removeNode(String nodeId) {
        nodes.remove(nodeId);
    }
    
    public Set<ClusterNode> getAliveNodes() {
        return nodes.values().stream()
                .filter(node -> node.getStatus() == NodeStatus.ALIVE)
                .collect(Collectors.toSet());
    }
    
    public void updateNodeStatus(String nodeId, NodeStatus status) {
        ClusterNode node = nodes.get(nodeId);
        if (node != null) {
            node.setStatus(status);
            node.updateHeartbeat();
        }
    }
    
    /**
     * 随机获取指定数量的存活节点
     * @param count 需要获取的节点数量
     * @return 随机选择的存活节点列表
     */
    public List<ClusterNode> getRandomPeers(int count) {
        Set<ClusterNode> aliveNodes = getAliveNodes();
        
        // 如果存活节点数量小于等于请求数量，直接返回所有存活节点
        if (aliveNodes.size() <= count) {
            return new ArrayList<>(aliveNodes);
        }
        
        // 随机选择指定数量的节点
        List<ClusterNode> allNodes = new ArrayList<>(aliveNodes);
        List<ClusterNode> selectedNodes = new ArrayList<>();
        Random random = new Random();
        
        for (int i = 0; i < count && !allNodes.isEmpty(); i++) {
            int randomIndex = random.nextInt(allNodes.size());
            selectedNodes.add(allNodes.remove(randomIndex));
        }
        
        return selectedNodes;
    }

    public void handleNodeJoin(GossipMessage message) {
        ClusterNode node = JSON.parseObject(JSON.toJSONString(message.getPayload()), ClusterNode.class);
        addNode(node);
    }

    public void handleNodeLeave(GossipMessage message) {
        String nodeId = (String) message.getPayload();
        removeNode(nodeId);
    }

    public void handleStatusUpdate(GossipMessage message) {
        ClusterNode node = JSON.parseObject(JSON.toJSONString(message.getPayload()), ClusterNode.class);
        updateNodeStatus(node.getNodeId(), node.getStatus());
    }
    
    /**
     * 处理数据消息
     * 1. 持久化消息到存储
     * 2. 转发消息给其他节点
     */
    private void handleDataMessage(GossipMessage gossipMessage) {
        try {
            // 从 payload 中提取消息数据
            DataMessagePayload payload = (DataMessagePayload) gossipMessage.getPayload();
            String topic = payload.getTopic();
            
            // 持久化消息
            boolean stored = storage.store(topic, payload);
            
            if (stored) {
                log.info("Successfully stored message {} for topic {}", 
                    gossipMessage.getMessageId(), topic);
                
                // 转发消息给其他节点
                gossipManager.spreadMessage(gossipMessage);
            } else {
                log.error("Failed to store message {} for topic {}", 
                    gossipMessage.getMessageId(), topic);
            }
            
        } catch (Exception e) {
            log.error("Error handling data message: {}", e.getMessage(), e);
        }
    }
} 